#include <opencv2/opencv.hpp>  
#include <complex>

using namespace cv;
using namespace std;

void fft(Mat img, Mat& result) {

  // 获取图像的宽、高
  int M = img.rows;
  int N = img.cols;

  // 将图像扩展到最佳尺寸  
  M = getOptimalDFTSize(M);
  N = getOptimalDFTSize(N);

  // 将图像补零到最佳尺寸
  Mat padded;                           
  copyMakeBorder(img, padded, 0, M - img.rows, 0, N - img.cols, BORDER_CONSTANT, Scalar::all(0));

  // 定义实部、虚部矩阵
  Mat planes[] = {Mat_<float>(padded), Mat::zeros(padded.size(), CV_32F)};
  Mat complexImg;
  merge(planes, 2, complexImg);    

  // 进行离散傅里叶变换
  dct(complexImg, complexImg);           
    
  // 将频谱移到中心位置
//   Mat tmp;
//   copyTo(complexImg, tmp, complexImg(Rect(0, 0, N/2, M/2)));
//   copyTo(complexImg, tmp, complexImg(Rect(N/2, 0, N/2, M/2)));
//   copyTo(complexImg, tmp, complexImg(Rect(0, M/2, N/2, M/2)));
//   copyTo(complexImg, tmp, complexImg(Rect(N/2, M/2, N/2, M/2)));
    Mat q0(complexImg, Rect(0, 0, N/2, N/2));
    Mat q1(complexImg, Rect(N/2, 0, N/2, M/2));
    Mat q2(complexImg, Rect(0, M/2, N/2, M/2)); 
    Mat q3(complexImg, Rect(N/2, M/2, N/2, M/2));

    q0.copyTo(complexImg(Rect(N/2, M/2, N/2, M/2)));
    q1.copyTo(complexImg(Rect(0, M/2, N/2, M/2)));  
    q2.copyTo(complexImg(Rect(N/2, 0, N/2, M/2)));
    q3.copyTo(complexImg(Rect(0, 0, N/2, M/2)));

  // 将复数转化为幅值,即傅里叶变换的模
  split(complexImg, planes);                  
  magnitude(planes[0], planes[1], result);

  // 缩放到0-1之间 then convert to 8bit
  result += Scalar::all(1);
  log(result, result);
  normalize(result, result, 0, 1, NORM_MINMAX); 
  result.convertTo(result, CV_8UC1);
}

int main() {
  
  Mat img = imread("/home/sunfy/sim_code2023_circle/src/picture/depth.png", 0);
  Mat result;
  fft(img, result);
  imshow("result", result);
  waitKey();

  return 0;
}